import pandas as pd
from fuzzywuzzy import process
from fuzzywuzzy import fuzz


uniprot_file = 'uniprot.xlsx'
tcmsp_file = 'tcmsp.xlsx'


def split_protein_names(names_str):
    names_list = names_str.split(' (')
    for index, value in enumerate(names_list):
        if value[-1] == ')':
            names_list[index] = value[:-1]
    return names_list


def main():
    uniprot_df = pd.read_excel(uniprot_file, usecols=['Protein names', 'Gene names  (primary )'])
    tcmsp_df = pd.read_excel(tcmsp_file, usecols=['Target name'])

    db_proteins = list()
    for index, value in uniprot_df['Protein names'].items():
        db_proteins.append(split_protein_names(value))

    result_dict = {"Target name": [], "Gene names  (primary )": []}
    for i, vi in tcmsp_df['Target name'].items():
        result_dict["Target name"].append(vi)
        result_dict["Gene names  (primary )"].append("")
        for j, vj in enumerate(db_proteins):
            result = process.extract(vi, vj, scorer=fuzz.token_sort_ratio)
            if result[0][1] == 100:
                if len(result_dict["Gene names  (primary )"][i]) != 0:
                    result_dict["Gene names  (primary )"][i] += "/"
                result_dict["Gene names  (primary )"][i] += uniprot_df.loc[j, "Gene names  (primary )"]
                print(i, vi, result, result_dict["Gene names  (primary )"][i])
    result_df = pd.DataFrame(result_dict)
    result_df.to_excel("result.xlsx", index=False)


if __name__ == "__main__":
    main()
